import numpy as np

from htm.Connection import Connection


class SequentialTMRegion:

    def __init__(self, connParams):
        # Parameters
        self.basalConnection = Connection(connParams)
        self.learn = True
        # Initializations
        self.monitorInfo = MonitorInfo()
        # Temporal States
        self.activeCellIds = np.array([])
        self.winnerCellIds = np.array([])
        self.predictedCellIds = np.array([])
        self.predictedActiveCellIds = np.array([])
        self.activeSegments = np.array([])
        self.matchingSegments = np.array([])
        self.potentialOverlaps = np.array([])

    def getNumOfCells(self):
        return self.basalConnection.columnCount * self.basalConnection.cellsPerColumn

    def compute(self, feedforwardInput, basalInput, basalGrowthCandidate):
        feedforwardInput.sort(), basalInput.sort(), basalGrowthCandidate.sort()
        self.depolarize(basalInput)
        self.activateCells(feedforwardInput, basalInput, basalGrowthCandidate)

    def depolarize(self, basalInput):
        connectedOverlaps, potentialOverlaps = self.basalConnection.computeActivity(basalInput)
        activeSegments = self.basalConnection.getActiveSegments(connectedOverlaps)
        matchingSegments = self.basalConnection.getMatchingSegments(potentialOverlaps)
        cellsForActiveSegments = self.basalConnection.mapSegmentsToCells(activeSegments)
        self.predictedCellIds = np.unique(cellsForActiveSegments)
        self.activeSegments = activeSegments
        self.matchingSegments = matchingSegments
        self.potentialOverlaps = potentialOverlaps

    def activateCells(self, feedforwardInput, basalInput, basalGrowthCandidate):
        cellsPerColumn = self.basalConnection.cellsPerColumn
        self.monitorInfo.setStatisticInfo(feedforwardInput, self.predictedCellIds // cellsPerColumn)
        # Calculate active cells
        correctPredictedCellIds = self.predictedCellIds[np.in1d(self.predictedCellIds // cellsPerColumn, feedforwardInput)]
        burstingColumns = np.setdiff1d(feedforwardInput, self.predictedCellIds // cellsPerColumn)
        cellsInBurstingColumns = Connection.getCellsInColumns(burstingColumns, cellsPerColumn)
        newActiveCells = np.union1d(correctPredictedCellIds, cellsInBurstingColumns)
        # Calculate basal learning
        learningActiveSegments = self.basalConnection.filterSegmentsByCell(self.activeSegments, correctPredictedCellIds)
        cellsForMatchingSegments = self.basalConnection.mapSegmentsToCells(self.matchingSegments)
        matchingCells = np.unique(cellsForMatchingSegments)
        matchingCellsInBurstingColumns = matchingCells[np.in1d(matchingCells // cellsPerColumn, burstingColumns)]
        learningMatchingSegments = self.basalConnection.chooseBestSegmentPerColumn(matchingCellsInBurstingColumns, self.matchingSegments, self.potentialOverlaps)
        learningMatchingCells = self.basalConnection.mapSegmentsToCells(learningMatchingSegments)
        burstingColumnsWithNoMatch = np.setdiff1d(burstingColumns, matchingCells // cellsPerColumn)
        newSegmentCells = self.basalConnection.getCellsWithFewestSegments(burstingColumnsWithNoMatch)
        learningCells = np.union1d(np.union1d(correctPredictedCellIds, learningMatchingCells), newSegmentCells)
        correctMatchingMask = np.in1d(cellsForMatchingSegments // cellsPerColumn, feedforwardInput)
        segmentsToPunish = self.matchingSegments[~correctMatchingMask]
        # Learn
        if self.learn:
            for learningSegments in (learningActiveSegments, learningMatchingSegments):
                self.basalConnection.learnOnExisting(learningSegments, basalInput, basalGrowthCandidate)
            self.basalConnection.adjustActiveSynapses(segmentsToPunish, basalInput, -self.basalConnection.permanencePunishment, 0)
            self.basalConnection.learnOnNew(newSegmentCells, basalGrowthCandidate)
        # For Output
        self.activeCellIds = newActiveCells
        self.winnerCellIds = learningCells
        self.predictedActiveCellIds = correctPredictedCellIds


class MonitorInfo:

    def __init__(self):
        self.predictedActiveColumnsList = []
        self.predictedInactiveColumnsList = []

    def setStatisticInfo(self, feedforwardInput, predictedColumnIds):
        predictedActiveColumns = set()
        predictedInactiveColumns = set()
        for pco in predictedColumnIds:
            if pco in feedforwardInput:
                predictedActiveColumns.add(pco)
            else:
                predictedInactiveColumns.add(pco)
        self.predictedActiveColumnsList.append(predictedActiveColumns)
        self.predictedInactiveColumnsList.append(predictedInactiveColumns)
